from langchain.agents import create_react_agent, AgentExecutor
from langchain_community.agent_toolkits.load_tools import load_tools
from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_core.tools import tool, Tool
from langchain_community.chat_message_histories import ChatMessageHistory

llm = Tongyi()


def tool1(user_input):
    return '小米手机非常好用，得到了广泛的好评，值得推荐，用网络热词《YYDS》形容也不为过。'


@tool
def tool2(user_input):
    """
    我曾经咨询过类似的问题
    :param user_input:
    :return:
    """
    return '我曾经咨询过类似的问题'


tools = [
    Tool(func=tool1, name='tool1', description="当咨询小米手机时，使用这个工具，只返回工具的结果"),
    tool2,
]

template = '''A请尽可能地回答以下问题。您可以使用以下工具:

           {tools}

           使用以下格式:

           Question: the input question you must answer
           Thought: you should always think about what to do
           Action: the action to take, should be one of [{tool_names}]
           Action Input: the input to the action
           Observation: the result of the action
           ... (this Thought/Action/Action Input/Observation can repeat N times)
           Thought: I now know the final answer
           Final Answer: the final answer to the original input question

           Begin!

           Question: {input}
           Thought:{agent_scratchpad}'''

prompt = ChatPromptTemplate.from_template(template)

agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

history = ChatMessageHistory()


def get_history():
    return history


runnable_history = RunnableWithMessageHistory(
    agent_executor,
    get_history,
    input_messages_key="input",
    history_messages_key="history",
)


def get_tool_agent(user_input):
    res = runnable_history.invoke({'input': user_input})
    print(res)
    return res